package com.yatop.lalmbda.ai;

import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisOutput;
import com.alibaba.dashscope.aigc.imagesynthesis.ImageSynthesisResult;
import com.alibaba.dashscope.utils.JsonUtils;
import com.yatop.lalmbda.ai.model.*;
import com.yuyaogc.lowcode.engine.context.Context;
import com.yuyaogc.lowcode.engine.plugin.activerecord.Criteria;
import com.yuyaogc.lowcode.engine.plugin.activerecord.Db;
import com.yuyaogc.lowcode.engine.plugin.activerecord.wrapper.LambdaQueryWrapper;
import com.yuyaogc.lowcode.engine.util.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Timestamp;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.*;
public class TaskScheduler {
    public static final Logger log = LoggerFactory.getLogger(TaskScheduler.class);

    private static final int QUEUE_CAPACITY = Integer.MAX_VALUE; // 队列的容量
    private final BlockingQueue<Long> taskQueue = new LinkedBlockingQueue<>(QUEUE_CAPACITY);
    private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
    private final ExecutorService threadPool = Executors.newCachedThreadPool();

    private  Context context;

    public void startScheduler() {
        // 每隔 5 秒钟检查队列并尽可能处理所有任务
        scheduler.scheduleAtFixedRate(() -> {
            processAllAvailableTasks();
        }, 0, 5, TimeUnit.SECONDS);
    }

    private void processAllAvailableTasks() {
        // 将队列中的所有任务移到一个临时列表中
        Queue<Long> tempQueue = new LinkedList<>();
        taskQueue.drainTo(tempQueue);

        // 提交所有任务到线程池
        tempQueue.forEach(this::processTask);
        System.out.println("Processed " + tempQueue.size() + " tasks");
    }

    public void addTask(Long task, Context context) {
        try {
            // 如果队列满，阻塞等待
            this.context = context;
            taskQueue.put(task);
            System.out.println("Task added: " + task);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            System.err.println("Error adding task: " + e.getMessage());
        }
    }

    private void processTask(Long task) {
        threadPool.submit(() -> {
            // 模拟任务处理逻辑
            try {

                System.err.println(context.getApp() + "================ ");
                try (Context context1 = new Context(context.getUserId(), Db.getConfig())){
                    Account account =  Account.selectOne(Account.class, Criteria.equal("userId", context.getUserId()));
                    if(null == account){
                       //TODO  throw new EngineException("当前账户为空");
                    }


                    Task task1  = Task.selectOne(Task.class, new LambdaQueryWrapper<Task>().eq(Task::getId,task));

                    String model = (String) task1.get("model");

                    AiModel aiModel = null;
                    if(null!=model){
                        aiModel =   AiModel.selectOne(AiModel.class ,new LambdaQueryWrapper<AiModel>().eq(AiModel::getModel,model));
                    }


                    String promptContext = aiModel.getName();
                    String size = aiModel.getSize();
                    String apiKey = account.getApiKey();
                    String modelName = aiModel.getModel();

                    // 如果有自定义的prompt，则使用自定义的prompt
                    if( StringUtils.isNotBlank(task1.getTaskName())){
                        promptContext = task1.getTaskName();
                    }

                    String taskId =  Util.createAsyncImage(promptContext, size, modelName, apiKey);
                    ImageSynthesisResult result = Util.waitAsyncTask(taskId, account.getApiKey());


                    ImageSynthesisOutput output = result.getOutput();
                    if(StringUtils.equals(output.getTaskStatus(), "SUCCEEDED")){
                        task1.set("model", aiModel.getId());
                        task1.setRequestId(result.getRequestId());
                        task1.setTaskId(taskId);
                        task1.setTaskStatus("SUCCEEDED");
                        task1.setTaskContext(JsonUtils.toJson(result));
                        task1.setUpdateTime(new Timestamp(System.currentTimeMillis()));
                        task1.setTaskEndTime(new Timestamp(System.currentTimeMillis()));
                        task1.update();


                        List<Map<String,String>> datas = output.getResults();
                        for(Map<String,String> data :datas){
                            Product product = new Product();
                            product.setUserId(context.getUserId());
                            product.setUrl(data.get("url"));
                            product.setOrigPrompt(data.get("orig_prompt"));
                            product.setCreateTime(new Timestamp(System.currentTimeMillis()));
                            product.setUpdateTime(new Timestamp(System.currentTimeMillis()));
                            product.save();
                        }
                    }


                    TaskDetail taskDetail = new TaskDetail();
                    taskDetail.set("task", task);
                    taskDetail.setTaskContext(JsonUtils.toJson(result));
                    taskDetail.setTaskId(taskId);
                    taskDetail.setUserId(context.getUserId());
                    taskDetail.setRequestId(null);
                    taskDetail.setCreateTime(new Timestamp(System.currentTimeMillis()));
                    taskDetail.setTaskType("jpg");
                    taskDetail.setTaskName(task1.getTaskName());
                    taskDetail.setTaskBeginTime(new Timestamp(System.currentTimeMillis()));
                    taskDetail.setTaskEndTime(new Timestamp(System.currentTimeMillis()));
                    taskDetail.setUpdateTime(new Timestamp(System.currentTimeMillis()));
                    taskDetail.save();

                    System.err.println(task1.getTaskName() + " ============");
                }

                //Thread.sleep(1000); // 模拟任务执行时间
                System.out.println("Task processed: " + task);
            } catch (Exception e) {
                log.error("========== {}", e);
                //TODO 错误日志处理
                Thread.currentThread().interrupt();
                e.printStackTrace();
                System.err.println("Task interrupted: " + task);
            }
        });
    }

    public void stopScheduler() {
        scheduler.shutdown();
        threadPool.shutdown();
    }

    public static void main(String[] args) {
        TaskScheduler scheduler = new TaskScheduler();
        scheduler.startScheduler();

        // 添加一些任务到队列
        for (int i = 1; i <= 10; i++) {
            // scheduler.addTask("Task " + i);
        }

        // 为了保持主线程运行，防止程序立即结束
        try {
            Thread.sleep(10000); // 等待 10 秒钟
        } catch (InterruptedException e) {
            e.printStackTrace();
        }

        //scheduler.addTask("Task " + 11);

        //scheduler.stopScheduler();
    }
}